package com.zcyr.util;

import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.parser.ParserException;
import com.zcyr.core.sql.OGSQLStatementParser;
import com.zcyr.core.sql.statement.OGDropSchemaStatement;
import com.zcyr.core.sql.visitor.OGOutputVisitor;
import com.zcyr.core.sql.visitor.OGTrimVisitor;
import com.zcyr.pojo.SQLResult;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.List;

public class SQLTrim {
    private static final Logger logger = LoggerFactory.getLogger(SQLTrim.class);

    private static void handleSQLExprs(List<SQLExpr> exprs, String suffix){
        for (SQLExpr expr : exprs){
            handleSQLExpr(expr, suffix);
        }
    }

    private static void handleSQLExpr(SQLExpr expr, String suffix){
        String name;
        if (expr instanceof SQLIdentifierExpr) {
            name = ((SQLIdentifierExpr) expr).getName();
            if (!name.equals("")) {
                ((SQLIdentifierExpr) expr).setName(name + suffix);
            }
        } else {
            throw new ParserException(String.format("不支持该'%s'类型的Expr处理...", expr.getClass().getSimpleName()));
        }
    }

    private static String handleSQL(String sql, String suffix){
        SQLStatement statement = new OGSQLStatementParser(sql).parseStatement();
        if (statement instanceof SQLGrantStatement) {
            SQLGrantStatement stmt = (SQLGrantStatement) statement;
            handleSQLExpr(stmt.getTo(), suffix);

        } else if (statement instanceof SQLRevokeStatement) {
            SQLRevokeStatement stmt = (SQLRevokeStatement) statement;
            handleSQLExpr(stmt.getFrom(), suffix);

        } else if (statement instanceof SQLDropTableStatement){
            SQLDropTableStatement stmt = (SQLDropTableStatement) statement;
            List<SQLExprTableSource> exprs = stmt.getTableSources();
            for (SQLExprTableSource expr : exprs){
                String schema = expr.getSchema();
                if (schema != null && !schema.equals("")){
                    expr.setSchema(schema + suffix);
                }
            }

        } else if(statement instanceof SQLCreateTableStatement) {
            SQLCreateTableStatement stmt = (SQLCreateTableStatement) statement;
            String schema = stmt.getSchema();
            if (schema != null && !schema.equals("")){
                stmt.setSchema(schema + suffix);
            }

        } else if (statement instanceof SQLDropUserStatement) {
            SQLDropUserStatement stmt = (SQLDropUserStatement) statement;
            handleSQLExprs(stmt.getUsers(), suffix);

        } else if (statement instanceof OGDropSchemaStatement) {
            OGDropSchemaStatement stmt = (OGDropSchemaStatement) statement;
            stmt.getSchemas().forEach(schema -> handleSQLExpr(schema, suffix));

        } else if (statement instanceof SQLCreateUserStatement) {
            SQLCreateUserStatement stmt = (SQLCreateUserStatement) statement;
            handleSQLExpr(stmt.getUser(), suffix);

        }else {
            logger.info(String.format("不支持该'%s'类型的SQLStatement隔离处理,使用OGTrimVisitor, SQL: %s", statement.getClass().getSimpleName(), sql));
            statement.accept(new OGTrimVisitor(suffix));
        }
        StringBuilder out = new StringBuilder();
        statement.accept(new OGOutputVisitor(out));
        System.out.println(out);
        return out.toString();
    }

    /**
     * 使用druid的sqlParser解析sql语句, 并在正确的地方添加后缀, 以实现用户与用户的sql隔离.
     * 如果sqlParser不兼容postgresql的某些语法导致解析失败, 则使用addSuffix处理.
     *
     * @param sql  要包装的sql
     * @param uuid 环境的uuid
     * @return 修改后的sql
     */
    public static String trimSQL(String sql, String uuid) {
        String[] sqlArray = sql.split(";");
        StringBuilder SQL = new StringBuilder();
        for (String s : sqlArray) {
            try {
                SQL.append(handleSQL(s, generateSuffix(uuid)));
            } catch (ParserException e) {
                logger.warn("SQLParse无法解析该SQL语句: " + sql + " ParserException: " + e.getMessage());
                SQL.append(s);
                //                throw new RuntimeException(e);
            } catch (Exception e){
                logger.error("SQLParse解析异常 SQL: " + s, e);
                SQL.append(s);
            }
            SQL.append(";");
        }
        return SQL.toString();
    }

//     /**
//     * 为用户执行的SQL语句做简单的修改, 以实现用户与用户的sql隔离
//     *
//     * @param sql  要修改的sql
//     * @param uuid 使用的环境uuid
//     * @return 修改后的sql
//     */
//    public static String addSuffix(String sql, String uuid) {
//        //将字符串进行分割
//        String[] keywords = sql.split("\\s+");
//        int i = 0;
//        //用于保存最后返回的字符串
//        //找到USER的位置，其下一个数组元素，就是username,在username后面添加 "$" + uuid;
//        StringBuilder sqlBuilder = new StringBuilder();
//        for (String keyword : keywords) {
//            if (isRequiredSuffixKeyword(keyword)) {
//                if (isNotCreatableName(keywords[i + 1])) {
//                    continue;
//                }
//                keywords[i + 1] = keywords[i + 1] + generateSuffix(uuid);
//            }
//            sqlBuilder.append(" ").append(keyword);
//            i++;
//        }
//        sql = sqlBuilder.toString();
//        return sql;
//    }

    public static String removeSuffix(String target, String uuid) {
        return target.replace(generateSuffix(uuid), "");
    }

    public static void removeColumnSuffix(String colName, SQLResult result, String uuid) {
        for (HashMap<String, String> data : result.tableData) {
            data.put(colName, removeSuffix(data.get(colName), uuid));
        }
    }

    public static String generateSuffix(String uuid) {
        return "$" + uuid;
    }
}
